from fastapi import Depends, UploadFile
import typing as T
from datetime import datetime
import uuid
from milvus_model.hybrid.mgte import MGTEEmbeddingFunction
from sqlalchemy import select, update, desc


from ...pojo.dto import (
    UserKnowledgeBaseDTO,
    KnowledgeHeaderDTO,
    EmbeddingSearchParamsDTO,
    UpdateImageInfoDTO,
)
from ...pojo.entity import UserKnowledgeBase, UserKnowledge, UserKnowledgeBaseFile
from ...pojo.vo import ImageInfoVO, UserKnowledgeBaseVO
from ...common.utils import loader_spliter
from ...init.Global import EmbeddingFunction
from ..mapper.UserKnowledgeBaseMapper import (
    _UserKnowledgeBaseMapper,
    UserKnowledgeBaseMapper,
)


class UserKnowledgeBaseServiceImpl:
    mapper: _UserKnowledgeBaseMapper = None
    embeddingFunction: MGTEEmbeddingFunction = None

    def __call__(
        self, mapper: UserKnowledgeBaseMapper, embeddingFunction: EmbeddingFunction
    ):
        self.embeddingFunction = embeddingFunction
        self.mapper = mapper
        return self

    async def create_knowledge_base(self, userKnowledgeBaseDTO: UserKnowledgeBaseDTO):
        createTime = datetime.now()
        alias = userKnowledgeBaseDTO.alias or f"{uuid.uuid4()}"
        entity = UserKnowledgeBase(
            alias=alias, createTime=createTime, userId=userKnowledgeBaseDTO.userId
        )
        await self.mapper.insert_by_entity(entity)

    async def get_knowledge_base_list_by_user_id(self, userId: int):
        sql = (
            select(UserKnowledgeBase)
            .where(UserKnowledgeBase.userId == userId)
            .order_by(desc(UserKnowledgeBase.createTime))
        )
        result = await self.mapper.select(sql)

        return [UserKnowledgeBaseVO(**item.model_dump()) for item in result]

    async def update_knowledge_base_name_by_knowledge_base_id(
        self, knowledgeBaseId: int, name: str
    ):
        sql = (
            update(UserKnowledgeBase)
            .where(UserKnowledgeBase.id == knowledgeBaseId)
            .values(alias=name)
        )
        await self.mapper.update(sql)

    async def delete_knowledge_base_by_knowledge_base_id(self, knowledgeBaseId: int):
        await self.mapper.delete_knowledgebase_by_id(knowledgeBaseId)

    async def upload(
        self,
        file: UploadFile,
        knowledgeHeader: KnowledgeHeaderDTO,
    ):
        # TODO 优化表结构，现在查询太麻烦
        file_name = file.filename
        type = file.content_type.split("/")[1]
        _uuid = f"{uuid.uuid4()}"
        params = dict(
            blob=await file.read(), file_name=file_name, file_type=type, uuid=_uuid
        )
        if type in ["png", "jpg", "jpeg", "jpeg"]:
            object_name = f"{_uuid}.{type}"
            params.update(
                object_name=object_name,
            )
        else:
            params.update(
                chunk_size=knowledgeHeader.chunk_size,
                chunk_overlap=knowledgeHeader.chunk_overlap,
                separators=knowledgeHeader.separators or "\n\n",
            )
        docs = loader_spliter(**params)
        contents = [doc.page_content for doc in docs]
        metadatas = [doc.metadata for doc in docs]
        embeddings = self.embeddingFunction.encode_queries(contents)
        await self.mapper.vector_insert_batch(
            [
                UserKnowledge(
                    knowledgeBaseId=knowledgeHeader.knowledgeBaseId,
                    userId=knowledgeHeader.userId,
                    content=content,
                    metadata=metadata,
                    embedding=embedding.detach().cpu().numpy(),
                ).model_dump()
                for content, metadata, embedding in zip(
                    contents, metadatas, embeddings.get("dense")
                )
            ]
        )
        await self.mapper.insert_by_entity(
            UserKnowledgeBaseFile(
                knowledgeBaseId=knowledgeHeader.knowledgeBaseId,
                fileName=file_name,
                alias=_uuid,
                type=metadatas[0].get("type"),
            )
        )
        if metadatas[0].get("type") == "image":
            return contents[0]

    async def search(self, embeddingSearchParamsDTO: EmbeddingSearchParamsDTO):
        embedding = (
            self.embeddingFunction.encode_queries([embeddingSearchParamsDTO.query])
            .get("dense")[0]
            .detach()
            .cpu()
            .numpy()
        )
        exp = f"knowledgeBaseId == {embeddingSearchParamsDTO.knowledgeBaseId}"
        if embeddingSearchParamsDTO.type == "image":
            exp = f'{exp} and metadata["type"] == "image"'
        else:
            exp = f'{exp} and metadata["type"] != "image"'
        know = await self.mapper.vector_search(
            embedding,
            exp,
            embeddingSearchParamsDTO.top_n,
        )
        know = sorted(know, key=lambda x: x.get("distance"), reverse=True)
        result = []
        for data in know:
            entity = data.get("entity")
            e = {
                "content": entity.get("content"),
                "type": entity.get("metadata").get("type"),
            }
            if e["type"] == "image":
                e["url"] = entity.get("metadata").get("url")
            if embeddingSearchParamsDTO.returnId:
                e["id"] = entity.get("metadata").get("uuid")
            result.append(e)
        return result

    async def get_images(self, knowledgeBaseId: int, count: int):
        exp = f"knowledgeBaseId == {knowledgeBaseId} and metadata['type'] == 'image'"
        know = await self.mapper.vector_query_by_expression(exp, count)
        result = []
        for data in know:
            metadata = data.get("metadata")
            result.append(
                ImageInfoVO(
                    url=metadata.get("url"),
                    id=metadata.get("uuid"),
                    fileName=metadata.get("file_name"),
                    content=data.get("content"),
                )
            )
        return result

    async def update_image_info(self, imageInfo: UpdateImageInfoDTO):
        # TODO 用upsert优化
        exp = f"knowledgeBaseId == {imageInfo.knowledgeBaseId} and metadata['uuid'] == '{imageInfo.id}'"
        entity: dict = (await self.mapper.vector_query_by_expression(exp, 1))[0]
        if not entity:
            return
        id = entity.pop("id")
        await self.mapper.vector_delete_by_id(id)
        embedding = (
            self.embeddingFunction.encode_queries([imageInfo.content])
            .get("dense")[0]
            .detach()
            .cpu()
            .numpy()
        )
        await self.mapper.vector_insert_batch(
            [
                UserKnowledge(
                    knowledgeBaseId=imageInfo.knowledgeBaseId,
                    userId=entity.get("userId"),
                    content=imageInfo.content,
                    metadata=entity.get("metadata"),
                    embedding=embedding,
                ).model_dump()
            ]
        )


UserKnowledgeBaseService = T.Annotated[
    UserKnowledgeBaseServiceImpl, Depends(UserKnowledgeBaseServiceImpl())
]
